/**
 * @file   Matrix.h
 * @author ubuntu <dadi@ubuntu>
 * @date   Mon May 31 23:07:38 2021
 * 
 * @brief  Define a class for Matrix computation.
 * 
 * 
 */

#ifndef __DADI_MATRIX__
#define __DADI_MATRIX__

#include<valarray>
#include<iostream>

#define TEMPLATE template <typename T>

TEMPLATE class Matrix;

TEMPLATE
class Matrix : public std::valarray<T>
{
private:
    size_t n_col;
    
public:
    /** 
     * Constructor. Create a matrix with m rows and n cols.
     * 
     * @param m Row size.
     * @param n Col size.
     * 
     * @return An m*n Matrix.
     */
    Matrix(size_t m, size_t n);
    /** 
     * Get the row size of the matrix.
     * 
     * 
     * @return 
     */
    size_t get_n_row() const;
    /** 
     * Get the col size of the matrix.
     * 
     * 
     * @return 
     */
    size_t get_n_col() const;
    /** 
     * Matrix(i,j) represents the ith row and jth col of the matrix.
     * 
     * 
     * @return 
     */
    T& operator()(size_t i, size_t j);
    const T& operator()(size_t i, size_t j) const;
    /** 
     * n_row*n_col matrix multiplies n_col*1 vector.
     * 
     * @param p 
     * 
     * @return 
     */
    std::valarray<T> operator*(const std::valarray<T> &p) const;
    /** 
     * Transpose matrix. 
     * 
     * 
     * @return 
     */
    Matrix<T> transpose() const;
};

TEMPLATE
Matrix<T>::Matrix(size_t m, size_t n) : std::valarray<T>(0.0, m * n), n_col(n)
{
};

TEMPLATE
size_t Matrix<T>::get_n_row() const
{
    if(n_col == 0)
	return 0;
    else
	return this->size() / n_col;
};

TEMPLATE
size_t Matrix<T>::get_n_col() const
{
    return n_col;
};

TEMPLATE
T& Matrix<T>::operator()(size_t i, size_t j)
{
    if((j >= 0) && (j < n_col))
	return (*this)[i * n_col + j];
    else
    {
	std::cerr << "The index is wrong." << std::endl;
	exit(-1);
    }
};

TEMPLATE
const T& Matrix<T>::operator()(size_t i, size_t j) const
{
    if((j >= 0) && (j < n_col))
	return (*this)[i * n_col + j];
    else
    {
	std::cerr << "The index is wrong." << std::endl;
	exit(-1);
    }
};

TEMPLATE
std::valarray<T> Matrix<T>::operator*(const std::valarray<T> &p) const
{
    size_t n_row = get_n_row();
    size_t n_col = get_n_col();
    if (n_col != p.size())
    {
	std::cerr << "Dimensions between the matrix and vector are not matching." << std::endl;
	exit(-1);
    }
    std::valarray<T> re(n_row);
    for(size_t i = 0; i < n_row; i++)
    {
	re[i] = 0;
	for(int j = 0; j < n_col; j++)
	    re[i] = re[i] + p[j] * (*this)(i,j);
    }
    return re;
};

TEMPLATE
Matrix<T> Matrix<T>::transpose() const
{
    size_t n_row = get_n_row();
    Matrix<T> T_Mat(n_col, n_row);
    for (size_t i = 0; i < n_col; i++)
        for (size_t j = 0; j < n_row; j++)
            T_Mat(i, j) = (*this)[j * n_col + i];
    return T_Mat;
}

#undef TEMPLATE

#else
//DO NOTHING
#endif




